#/usr/bin/python
# -*- coding: utf-8 -*-
import sys
sys.path.append('../../../icafe')
import parameters

from mesh.borderMesh import *
from copy import deepcopy
import scipy.sparse as sp

class Edge():
    '''
    This class defines an edge (part of the border) of a domain. It is a discrete (n-k)-manifold of the considered 
    domain (which is a n-manifold). An edge is indeed linked to a mesh of the domain from which it is extracted. 
    '''
    
    def __init__(self,UId, domainUId, mesh=None, interpmeshes=None, nodelistCId=None, nodelistUId=None, coord_expression=None, **kwargs):
        '''
        Edge constructor:
        @param domainUId: user identifier of the associated domain.
        @param interpMeshes: meshes of the edge without Lhs and Rhs computation method. Nodes of these meshes have the 
        same user identifier as corresponding node in associated meshes of associated domain.
        @param interpMeshMapToMesh: mapping of the interpolation meshes with meshes of the associated domain. 
        @param traceOperator: trace operator that extracts d.o.f. of the domain corresponding to the nodes of the edge.
        '''
        # related domain id
        self.UId = UId
        self.domainUId = domainUId
        self.interpMeshes = interpmeshes
        self.domainMesh_to_edgeMesh = {}
        self.edgeMesh_to_domainMesh = []
        self.nodesMap = {"TOMESHNODES":{}, }
        # calculated trace operator for the meshes of the domain
        
        self.traceOperator = [None]
        
#        args = self.__dict__ 
#        
#        for k in kwargs.keys():
#            if k in args.keys():
#                self.__setattr__(k,kwargs[k])
        
        if (nodelistCId or nodelistUId) and mesh:
            self.buildInterpMesh(mesh, nodelistCId=nodelistCId, nodelistUId=nodelistUId)
        elif pointlist and mesh:
            self.buildInterpMesh(mesh, pointlist=pointlist)
        elif coord_expression and mesh:
            self.buildInterpMesh(mesh, coord_expression=coord_expression)
#    def calcMarker(self, domain=None, nodeSet=None, elementSet=None, mesh=0, n_mesh=1 ):
#        '''
#        Returns a marker for the edge in order to compare it with other.
#        '''
#        pass
    def addMesh(self, mesh):
        '''
            @param mesh: Mesh.
        '''
        mid = len(self.interpMeshes)
        self.domainMesh_to_edgeMesh[mesh.UId]=mid
        
        self.buildInterpMesh(mesh)
        
    def removeMesh(self, meshUId):
        '''
            @param meshUId: str or int. User id of the mesh whose interpMesh must be removed.
        '''
        raise NotImplementedError  
    
    def getNodesFromMesh(self, mesh, output='CId'):
        '''
            
            @param meshUId: str or int. User id of the mesh.
        '''
        
        nodes = []
        
        try:
            nodes = self.nodesMap["TOMESHNODES"][mesh.UId]
            if output=='CId':
                nodes = [mesh.nodesMap['U2C'][ndUId] for ndUId in nodes]  
        except KeyError: 
            imesh = self.getInterpMesh(mesh.UId)   
            nnodes = len(imesh.nodesMap["C2U"])
            for nid in xrange(nnodes):
                if output == 'CId':
                    nodes.append(mesh.nodesMap['U2C'][imesh.nodesMap['C2U'][nid]])
                else:
                    nodes.append(imesh.nodesMap['C2U'][nid])
            self.nodesMap["TOMESHNODES"][mesh.UId]=nodes
        
        return nodes
    
    def getInterpMesh(self, meshUId):
        '''
            @param meshUId: str or int. User id of the mesh from which the interpMesh has been extracted.
        '''
        try:
            mid = self.domainMesh_to_edgeMesh[meshUId]
            return self.interpMeshes[mid]
        except IndexError:
            raise IndexError('Mesh '+str(meshUId)+' has not been added to the edge ', self.UId)
         
    def getTraceOperator(self, meshUId):
        '''
            @param meshUId: str or int. User identifier of the mesh from which to get the trace operator to the edge. 
        '''
        try:
            mid = self.domainMesh_to_edgeMesh[meshUId]
            return self.traceOperator[mid]
        except IndexError:
            raise IndexError('Mesh '+str(meshUId)+' has not been added to the edge ', self.UId)
             
    def buildTraceOperator(self, mesh):
        '''
           Compute the trace operator between a mesh and its corresponding interpMesh of the edge.
           @param mesh: Mesh. 
        '''
        try:
            mid = self.domainMesh_to_edgeMesh[mesh.UId]
        except ValueError:
            raise ValueError("Edge ", self.UId, " has nothing to do with mesh:", mesh.UId)
        
        imesh = self.interpMeshes[mid]
        nnodes = len(imesh.nodesMap["C2U"])
        
        if imesh.ndofs_per_node != mesh.ndofs_per_node:
            raise NotImplementedError('Mesh and edge have not the same number of degrees of freedom per nodes.')
        
        edofs = nnodes*imesh.ndofs_per_node
        mdofs = len(mesh.nodesMap["C2U"])*mesh.ndofs_per_node
        
        tOp = sp.lil_matrix((edofs,mdofs))
        for nid in xrange(nnodes):
            for dof in xrange(imesh.ndofs_per_node):
                try:
                    tOp[nid*imesh.ndofs_per_node+dof,mesh.nodesMap['U2C'][imesh.nodesMap['C2U'][nid]]*mesh.ndofs_per_node+dof]=1 
                except IndexError, arg:
                    print 'dof: ',dof, ', node: ', nid, ', node in mesh: ', mesh.nodesMap['U2C'][imesh.nodesMap['C2U'][nid]]
                    print 'row: ',nid*imesh.ndofs_per_node+dof
                    print 'col: ',imesh.nodesMap['C2U'][nid]*mesh.ndofs_per_node+dof
                    raise IndexError, arg
        try:
            self.traceOperator[mid] = tOp
        except IndexError:
            self.traceOperator.append(tOp)
            
        return tOp
    
    def buildInterpMesh(self, mesh, nodelistCId=None, nodelistUId=None, pointlist=None, coord_expression=None):
        

        imesh=InterpMesh()

           
        if nodelistCId:
            # nodes:
            nodes = deepcopy(mesh.getNodes(nodelistCId=nodelistCId))
            
            self.nodesMap["TOMESHNODES"][mesh.UId] = nodes["UId"]
            
            for i,ndcoo in enumerate(nodes["coord"]):
                imesh.addNode(nodes["UId"][i],ndcoo)
            
#            imesh.nodesCoo = [mesh.nodesCoo[nid] for nid in nodelistCId]
#            imesh.nodesMap["U2C"]=dict(zip(nodelistCId,range(len(nodelistCId))))
#            imesh.nodesMap["C2U"]=nodelistCId
#            for i,ndcoo in enumerate(imesh.nodesMap['C2COO']): 
#                imesh.nodesMap["COO2C"][mesh.nodesMap["C2COO"][nodelistCId[i]]]= i
#                imesh.nodesMap["C2COO"].append(ndcoo)
            
            # elements:
            
            for i,elems in enumerate([nodes["elements"][nid] for nid in xrange(len(nodes["UId"]))]):
#                nodes = mesh.getNodes(elemUId=elem)
                    
                elements = deepcopy(mesh.getElements(elemUId=elems))
                
                for i, type in enumerate(elements['types']):    
                    imesh.addElement(elements['UId'][i],elements['nodes'][i],type)
                    
#                    imesh.elementsNodes.append([mesh.nodesMap["C2U"][node] for node in nodes])
#                    imesh.elementsType.append(mesh.elementsType[mesh.elementsMap["U2C"][elem]])
#                    imesh.elementsMap["U2C"][elem]=nelem
#                    imesh.elementsMap["C2U"].append(elem)
#                    nelem+=1
            
            # enrichments:
            # TODO
             
            imesh.consolidateMesh()
        
        elif nodelistUId:
            
            # nodes:
            nodes = deepcopy(mesh.getNodes(nodelistUId=nodelistUId))
            
            self.nodesMap["TOMESHNODES"][mesh.UId] = nodes["UId"]
            
            for i,ndcoo in enumerate(nodes["coord"]):
                imesh.addNode(nodes["UId"][i],ndcoo)
            
#            imesh.nodesCoo = [mesh.nodesCoo[mesh.nodesMap["U2C"][nid]] for nid in nodelistUId]
#            imesh.nodesMap["U2C"]=dict(zip(nodelistUId,range(len(nodelistUId))))
#            imesh.nodesMap["C2U"]=[mesh.nodesMap["U2C"][nid] for nid in nodelistUId]
#            for i,ndcoo in enumerate(imesh.nodesMap['C2COO']): 
#                imesh.nodesMap["COO2C"][mesh.nodesMap["C2COO"][mesh.nodesMap["U2C"][nodelistUId[i]]]]= i
#                imesh.nodesMap["C2COO"].append(ndcoo)
            
            # elements:
            for elems in [nodes["elements"][nid] for nid in xrange(len(nodes["UId"]))]:
#                nodes = mesh.getNodes(elemUId=elem)
                    
                elements = deepcopy(mesh.getElements(elemUId=elems))
                
                for i, type in enumerate(elements['types']):    
                    
                    imesh.addElement(elements['UId'][i],elements['nodes'][i],type)
                    
                    
#            nelem=0
#            for i,elems in enumerate([mesh.nodesElements[mesh.nodesMap["U2C"][nid]] for nid in nodelistUId]):
#                for elem in elems:
#                    nodes = mesh.elementsNodes[mesh.elementsMap["U2C"][elem]]
#                    imesh.elementsNodes.append([mesh.nodesMap["C2U"][node] for node in nodes])
#                    imesh.elementsType.append(mesh.elementsType[mesh.elementsMap["U2C"][elem]])
#                    imesh.elementsMap["U2C"][elem]=nelem
#                    imesh.elementsMap["C2U"].append(elem)
#                    nelem+=1
            
            # enrichments:
            # TODO
            
            imesh.consolidateMesh()       
            
            
                
        elif pointlist:
            raise NotImplementedError
        elif coord_expression:
            raise NotImplementedError
        
        else:
            raise NotImplementedError
            # find nodes on edge and build a nodelist:
            
            # add nodes and degenerated elements:
            
            # consolidate the mesh:
            imesh.consolidateMesh()
        
        try:
            self.interpMeshes.append(imesh)
        except AttributeError:
            self.interpMeshes = [imesh]
        
        mid = len(self.interpMeshes)-1
        
        self.domainMesh_to_edgeMesh[mesh.UId]=mid
        self.edgeMesh_to_domainMesh.append(mesh.UId)
        
        self.buildTraceOperator(mesh)
        
        
    
class Point(Edge):
    '''
    Defines an edge as a point  
    '''
    
    def __init__(self, edge, node, mesh = 0, n_mesh = 1):
        '''
        
        '''
        kwargs = edge.__dict__()
        super().__init__(**kwargs)
        self.calcMarker()
        
    
    def calcMarker(self):
        self.marker = self.nodes[0]
    
class Curve(Edge):
    '''
    Defines a curve as an edge
    '''
    
    def __init__(self, edge, nodeset, mesh = 0, n_mesh = 1):
        '''
        
        '''
        super.__init__(self, nodeset, mesh, n_mesh)
        
    
    def calcMarker(self):
        pass
    
    
class Surface(Edge):
    '''
    Defines a surface as an edge
    '''
    
    def __init__(self, nodeset, mesh = 0, n_mesh = 1):
        '''
        
        '''
        super.__init__(self, nodeset, mesh, n_mesh)
    
    def calcMarker(self):
        pass
 
    
EDGE_CLASS = {
        "edge": Edge,
        "point": Point,          
        "curve": Curve,
        "surface": Surface,
        }

